# -*- coding: utf-8 -*-
"""
------------------------------------------------------------------------------
    File Name:  topics_and_transformations
    Author   :  wanwei1029
    Date     :  2019/3/15
    Desc     : 翻译：http://blog.geekidentity.com/nlp/gensim/gensim_2_topics-and-transformations_translation/
    上一章将批量文档，生成了字典，并通过字典生成了词袋模型，总的来说，词袋模型是一种很简单的模型，但我们可以在此基础
    上，生成更多更复杂且实用的新模型(模型也是需要训练数据生成的)，比如tfidf,lsi等等。
    注意：每个模型都有自己的作用，以及特定的输入和输出，比如tfidf模型，我们可以通过此模型得出一个向量中每个元素的权重
    信息，通过lsi模型可以生成主题。
------------------------------------------------------------------------------
"""
import logging
import os
from gensim import corpora, models, similarities
from pprint import pprint

TEMP_FOLDER = r"D:\nas\gensim\temp"
print('Folder "{}" will be used to save temporary dictionary and corpus.'.format(TEMP_FOLDER))
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)

# def tf_idf_model():


def lsi_model():
    raw_data_list = list()
    raw_data_list.append("王宝强 马蓉 离婚 马蓉 否认")
    raw_data_list.append("中国 美国 贸易 谈判")
    raw_data_list.append("王宝强 打官司")
    raw_data_list.append("王宝强 没钱 打官司")
    raw_data_list.append("马蓉 宋杰 在一起")
    raw_data_list.append("中国 美国 谈判 结局")
    raw_data_list.append("中国 美国 结束 谈判")
    raw_data_list.append("中国 美国 最终 条款")
    texts = [[word for word in raw_data.split(" ")] for raw_data in raw_data_list]
    pprint(texts)
    dictionary = corpora.Dictionary(texts)
    pprint(dictionary.token2id)
    bow_model = [dictionary.doc2bow(text) for text in texts]
    pprint(bow_model)
    # tfidfmodel的参数类型必须是model，但bow_model目前只是一个二维数组，类型不对，不能直接生成tfidf模型。
    # 纠正，可以直接生成，但是不是用tfidfmodel，而是用TfidfModel
    # 再次补充，模型调用要用[]，而不是(),要不就会报model is not callable.
    corpora.MmCorpus.serialize(os.path.join(TEMP_FOLDER, 'my_test.mm'), bow_model)
    bow_model = corpora.MmCorpus(os.path.join(TEMP_FOLDER, 'my_test.mm'))
    # tfidf_model = models.tfidfmodel(bow_model)
    tfidf_model = models.TfidfModel(bow_model)
    pprint(tfidf_model)



def demo():
    """
    """
    dictionary = corpora.Dictionary.load(os.path.join(TEMP_FOLDER, 'deerwester.dict'))
    corpus = corpora.MmCorpus(os.path.join(TEMP_FOLDER, 'deer_test.mm'))
    print(corpus)
    # 可以看出，bow模型corpus是由9篇文档生成，包含12个特征，28个非空元素。
    # 9个文档，很好理解，因为生成模型用到了9个文档
    # 12个特征，是指词典的大小，因为需要将每个文档转换成相同大小的向量，所以长度不够的文档会补齐，也就是说，最后生成的向量模型，是一个9*12的矩阵。
    # 28是指矩阵中一共有28个非空元素，所以矩阵的密码可以计算为：25.926（28/(9*12))
    tfidf = models.TfidfModel(corpus)
    # 将bow模型转换为tfidf模型。
    print(tfidf)
    # 用tfidf模型计算一个文本向量的tf-idf值。
    doc_bow = [(0, 1), (1, 1)]
    print(tfidf[doc_bow])
    # 对整个语料库计算tf-idf值。
    corpus_tfidf = tfidf[corpus]
    for doc in corpus_tfidf:
        print(doc)

    # 生成lsi模型
    lsi = models.LsiModel(corpus_tfidf, id2word=dictionary, num_topics=2)
    # 将原语料库转换成lsi
    lsi.print_topics(2)
    corpus_lsi = lsi[corpus_tfidf]
    for doc in corpus_lsi:
        print(doc)
    # 保存模型
    lsi.save(os.path.join(TEMP_FOLDER, 'model.lsi'))
    # 加载模型
    # lsi = models.LsiModel.load(os.path.join(TEMP_FOLDER, 'model.lsi'))


if __name__ == '__main__':
    test_method = "lsi_model"
    if test_method == "demo":
        demo()
    elif test_method == "lsi_model":
        lsi_model()
